# _*_ encoding:utf-8 _*_
import optparse
try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET
from generate_start_points import *
import math


def generate_vtype(v_types):
    return v_types[random.randint(0, len(v_types) - 1)]['id']


def load_vtypes(filename):
    vtypes = []
    for event, eml in ET.iterparse(filename):
        if event == 'end':
            if eml.tag == 'vType':
                vtype = eml.attrib
                if 'probability' not in vtype:
                    print "vtype file error, for each vType, you must specify the probability attrib and the sum of " \
                            "probabilities should be equal to 1"
                    exit(1)
                vtypes.append(vtype)
    return vtypes


def assign_vtype(vtypes):
    prob = []
    for vtype in vtypes:
        prob.append(float(vtype['probability']))

    probability = random.random()
    sum = 0
    for i in range(len(prob)):
        sum += prob[i]
        if sum > probability:
            return vtypes[i]['id']


def get_options(args=None):
    optparser = optparse.OptionParser()
    optparser.add_option('-r', "--route-file", dest="rou_file", help="the input rou xml")
    optparser.add_option('-v', "--vtypes", dest='v_file', help="the input vtypes xml")
    optparser.add_option("-o", "--output", dest="o_file", help="the output rou xml")
    optparser.add_option("-n", "--net", dest="net", help="the defined network")
    optparser.add_option("-p", "--probability", dest="probability", type="float", help="the probability for one vehicle to perform an emergency stop")
    (options, args) = optparser.parse_args(args=args)
    if not options.rou_file or not options.v_file or not options.o_file or not options.net:
        optparser.print_help()
        exit(0)
    if not options.probability:
        options.probability = 0
    return options


def load_net(filename):
    tree = ET.ElementTree(file=filename)
    edges = {}
    for node in tree.getroot():
        if node.tag == 'edge':
            edges.setdefault(node.attrib['id'], float(node[0].attrib['length']))
    return edges


def main(options=None):
    # vtypes_file = "../../data/1022.vtypes.add.xml"
    # route_xml = "../../data/10221515.rou.xml"
    print options.v_file
    print options.rou_file
    vtypes = load_vtypes(options.v_file)
    tree = ET.ElementTree(file=options.rou_file)
    edges = {}
    if options.net:
        edges = load_net(options.net)
    count = 0
    with open(options.o_file, "wb") as f:
        f.write('<?xml version="1.0" encoding="UTF-8"?>\n')
        f.write('<routes xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:noNamespaceSchemaLocation="http://sumo.dlr.de/xsd/routes_file.xsd">\n')
        for child in tree.getroot():
            if child.tag == 'vehicle':
                count += 1
                route = child[0].attrib['edges'].split(' ')
                depart = math.floor(float(child.attrib['depart']))
                id = child.attrib['id']
                types = assign_vtype(vtypes)
                if len(edges) != 0 and route[0] in edges and route[-1] in edges:
                    line = '\t<vehicle arrivalPos="random" depart="{0}" departLane="best" departPos="base" departSpeed="random" id="{1}" type="{2}">\n'.format(depart, id, types)
                    f.write(line)
                    f.write('\t\t<route edges="{0}"/>\n'.format(child[0].attrib['edges']))
                    proba = random.random()
                    if proba < options.probability:
                        # set emergency stop
                        e_id = random.randint(2, len(route)) - 1
                        e_id = route[e_id]
                        while edges[e_id] < 10:
                            e_id = random.randint(2, len(route)) - 1
                            e_id = route[e_id]
                        end_pos = random.random() * (edges[e_id] - 10) + 10
                        # start_pos = end_pos - 0.2
                        duration = random.randint(1, 30)
                        f.write('\t\t<stop lane="{0}_0" endPos="{1}" duration="{2}"/>\n'.format(e_id, end_pos, duration))
                        pass
                    f.write('\t</vehicle>\n')
        f.write('</routes>')

    print 'total : %d' %(count)

if __name__ == '__main__':
    main(get_options())

